#pragma once

#include "libmesh/linear_implicit_system.h"
#include "libmesh/dense_matrix.h"
#include "libmesh/dense_vector.h"
#include "Eigen/Eigen"
#include "libmesh/equation_systems.h"
#include "libmesh/dof_map.h"
#include "libmesh/elem.h"
#include "libmesh/sparse_matrix.h"

#include "DofMap.h"

using namespace Eigen;
using namespace libMesh;
using Matrix = Eigen::MatrixXd;
class EbESystem : public LinearImplicitSystem
{
public:
    EbESystem(EquationSystems& es,
        const std::string& name,
        const unsigned int number): LinearImplicitSystem(es, name, number)
    {}
    void add_matrix(DenseMatrix<Real>& Ke)
    {
        //Ke.m();
        //Ke.n();
        //MatrixXd A_local;
        //A_local.resize(Ke.m(), Ke.n());
        //for (size_t i = 0; i < Ke.m(); i++)
        //{
        //    for (size_t j = 0; j < Ke.n(); j++)
        //    {
        //        A_local(i, j) = Ke(i, j);
        //    }
        //}
        _A.push_back(Ke);

    }
    void add_rhs(DenseVector<Real>& Fe)
    {
        //VectorXd b_local(Fe.size());
        //for (size_t i = 0; i < Fe.size(); i++)
        //{
        //    b_local(i) = Fe(i);
        //}
    }

    virtual void solve() override;
private:
    std::vector<DenseMatrix<Real>> _A;
    std::vector<VectorXd> _b;
    void ebe_solve_cg();
    void ebe_solve_bicgstab();
    void A(int i)
    {
    }

    // y = Ax;
    void linear_operator(NumericVector<Number>& y, NumericVector<Number>& x)
    {
        const auto& es = get_equation_systems();
        const auto& mesh = es.get_mesh();
        const auto& dof_map = get_dof_map();
        std::vector<dof_id_type> dof_indices;
        y.zero();
        for (const auto& elem : mesh.active_local_element_ptr_range())
        {
            dof_map.dof_indices(elem, dof_indices);
            DenseVector<Real> be(dof_indices.size());
            for (size_t i = 0; i < dof_indices.size(); i++)
            {
                be(i) = x(dof_indices[i]);
            }
            auto id = elem->id();
            DenseVector<Real> temp;
            _A[id].vector_mult(temp, be);
            y.add_vector(temp , dof_indices);
        }
    }

    //r <- b - A * x
    void axpby(NumericVector<Number>& r, NumericVector<Number>& b, NumericVector<Number>& y)
    {
        assert(b.size() == y.size());
        assert(r.size() == y.size());
        for (size_t i = 0; i < b.size(); i++)
        {
            r.set(i, b(i) - y(i));
        }
    }

    void axpby(NumericVector<Number>& r, NumericVector<Number>& b, NumericVector<Number>& y, Real scale)
    {
        assert(r.size() == y.size());
        assert(b.size() == y.size());
        for (size_t i = 0; i < b.size(); i++)
        {
            r.set(i, b(i) + scale * y(i));
        }
    }

    Real dot(NumericVector<Number>& b, NumericVector<Number>& y)
    {
        assert(b.size() == y.size());
        Real square = 0;
        for (size_t i = 0; i < b.size(); i++)
        {
            square += b(i) * y(i);
        }
        return square;  
    }
    void precondition(NumericVector<Number>& z, NumericVector<Number>& r)
    {
        const auto &diag = get_vector("diag");
        assert(z.size() == r.size());
        for (size_t i = 0; i < r.size(); i++)
        {
            //z.set(i, r(i) / diag(i));
            z.set(i, r(i) / 1);
        }
    }
    void copy(NumericVector<Number>& dest, NumericVector<Number>& src)
    {
        assert(dest.size() == src.size());
        for (size_t i = 0; i < dest.size(); i++)
        {
            dest(0);
            auto a = src(i);
            dest.set(i, src(i));
        }
    }

    void zero(NumericVector<Number>& target)
    {
        target.zero();
    }
    void one(NumericVector<Number>& target)
    {
        for (size_t i = 0; i < target.size(); i++)
        {
            target.set(i, 1);
        }
    }
    void axpy(NumericVector<Number>& x, NumericVector<Number>& y, Real scale)
    {
        assert(x.size() == y.size());
        for (size_t i = 0; i < x.size(); i++)
        {
            x.set(i, x(i) + scale*y(i));
        }
    }
};